## 源数据集量太大，从中取样n%出来作为新数据集

import os
import random
import shutil

def select_and_copy_images_labels(source_dir, target_dir, percentage=0.9):
    """
    随机选择一定比例的图片和对应的标签文件复制到目标文件夹。

    :param source_dir: 原始数据集的父文件夹路径，其中包含 'images' 和 'labels' 子文件夹
    :param target_dir: 目标数据集的父文件夹路径，将包含 'images' 和 'labels' 子文件夹
    :param percentage: 要选择的比例（默认值为 0.9）
    """
    # 定义源文件夹和目标文件夹
    source_images_dir = os.path.join(source_dir, 'images')
    source_labels_dir = os.path.join(source_dir, 'labels')
    target_images_dir = os.path.join(target_dir, 'images')
    target_labels_dir = os.path.join(target_dir, 'labels')

    # 创建目标文件夹及其子文件夹
    os.makedirs(target_images_dir, exist_ok=True)
    os.makedirs(target_labels_dir, exist_ok=True)

    # 获取所有图片文件名（假设图片后缀为 jpg 或 png 等）
    image_files = [f for f in os.listdir(source_images_dir) if f.endswith(('.jpg', '.png', '.jpeg'))]
    
    # 按文件名随机打乱顺序
    random.shuffle(image_files)

    # 选择指定比例的文件
    selected_count = int(len(image_files) * percentage)
    selected_images = image_files[:selected_count]

    # 复制图片和对应的标签文件
    for image_file in selected_images:
        # 图片文件路径
        src_image_path = os.path.join(source_images_dir, image_file)
        tgt_image_path = os.path.join(target_images_dir, image_file)

        # 对应的标签文件路径
        label_file = os.path.splitext(image_file)[0] + '.txt'
        src_label_path = os.path.join(source_labels_dir, label_file)
        tgt_label_path = os.path.join(target_labels_dir, label_file)

        # 复制图片
        shutil.copy(src_image_path, tgt_image_path)
        
        # 如果标签文件存在则复制
        if os.path.exists(src_label_path):
            shutil.copy(src_label_path, tgt_label_path)

    print(f"已成功从 {source_dir} 中随机选择 {selected_count} 张图片及其标签文件，复制到 {target_dir}。")

source_folder = "D:/研究生/科研相关/无人艇专项/视觉数据集制作/total_202411251934"
target_folder = "D:/研究生/科研相关/无人艇专项/视觉数据集制作/new"
select_and_copy_images_labels(source_folder, target_folder, percentage=0.9)
